#ifndef TOOLS_MEMORY_POOL_H
#define TOOLS_MEMORY_POOL_H

#include <cstddef>
#include <mutex>
#include <vector>
#include <memory>
#include <cassert>

namespace Tool
{

template <typename T, size_t BlockSize = 2048>
class MemoryPool
{
public:
    // 类型定义
    using value_type = T;
    using pointer = T*;
    using reference = T&;
    using const_pointer = const T*;
    using const_reference = const T&;
    using size_type = size_t;
    using difference_type = ptrdiff_t;

    // 用于rebind到其他类型的内存池
    template <typename U>
    struct rebind {
        using other = MemoryPool<U, BlockSize>;
    };

    MemoryPool() noexcept
        : free_slots_(nullptr)
        , chunks_(nullptr)
    {
        static_assert(BlockSize > 0, "BlockSize must be greater than 0");
        expand();
    }

    MemoryPool(const MemoryPool& other) noexcept
        : MemoryPool()
    {

    }

    MemoryPool(MemoryPool&& other) noexcept
        : free_slots_(other.free_slots_)
        , chunks_(other.chunks_)
    {
        other.free_slots_ = nullptr;
        other.chunks_ = nullptr;
    }

    template <typename U>
    MemoryPool(const MemoryPool<U, BlockSize>& other) noexcept
        : MemoryPool()
    {

    }

    ~MemoryPool() noexcept
    {
        Chunk* current = chunks_;
        while (current != nullptr)
        {
            Chunk* next = current->next;
            delete current;
            current = next;
        }
    }

    MemoryPool& operator=(const MemoryPool& other) = delete;

    MemoryPool& operator=(MemoryPool&& other) noexcept
    {
        if (this != &other)
        {
            Chunk* current = chunks_;
            while (current != nullptr)
            {
                Chunk* next = current->next;
                delete current;
                current = next;
            }

            free_slots_ = other.free_slots_;
            chunks_ = other.chunks_;

            other.free_slots_ = nullptr;
            other.chunks_ = nullptr;
        }
        return *this;
    }

    void expand()
    {
        Chunk* new_chunk = new Chunk(BlockSize);
        new_chunk->next = chunks_;
        chunks_ = new_chunk;

        for (size_t i = 0; i < BlockSize - 1; ++i)
        {
            new_chunk->data[i].next = &new_chunk->data[i + 1];
        }
        new_chunk->data[BlockSize - 1].next = free_slots_;
        free_slots_ = new_chunk->data;
    }

    // 分配内存
    typename MemoryPool<T, BlockSize>::pointer
    allocate()
    {
        std::lock_guard<std::mutex> lock(mutex_);
        if (free_slots_ == nullptr)
        {
            expand();
        }

        pointer result = reinterpret_cast<pointer>(free_slots_);
        free_slots_ = free_slots_->next;
        return result;
    }

    void deallocate(pointer p)
    {
        std::lock_guard<std::mutex> lock(mutex_);

        Slot* slot = reinterpret_cast<Slot*>(p);
        slot->next = free_slots_;
        free_slots_ = slot;
    }

    template <typename... Args>
    void construct(pointer p, Args&&... args)
    {
        new (p) T(std::forward<Args>(args)...);
    }

    void destroy(pointer p)
    {
        p->~T();
    }

    size_type block_size() const noexcept { return BlockSize; }
    size_type object_size() const noexcept { return sizeof(T); }

private:
    union Slot
    {
        T element;
        Slot* next;
    };

    struct Chunk
    {
        Slot* data;
        Chunk* next;

        Chunk(size_type size) : next(nullptr)
        {
            data = reinterpret_cast<Slot*>(::operator new(size * sizeof(Slot)));
        }

        ~Chunk()
        {
            ::operator delete(data);
        }
    };

    Slot* free_slots_;
    Chunk* chunks_;
    std::mutex mutex_;
};

// 内存池分配器，用于标准容器
template<typename T, typename Pool>
class PoolAllocator
{
public:
    using value_type = T;

    PoolAllocator(Pool& pool) noexcept
        : m_pool(&pool)
    {
        static_assert(sizeof(T) <= pool.block_size(), "Type size exceeds memory pool block size");
    }

    template<typename U>
    PoolAllocator(const PoolAllocator<U, Pool>& other) noexcept
        : m_pool(other.m_pool)
    {
    }

    T* allocate(size_t n)
    {
        if (n != 1)
        {
            throw std::bad_alloc();
        }

        return static_cast<T*>(m_pool->allocate());
    }

    void deallocate(T* p, size_t n) noexcept
    {
        if (n != 1)
        {
            return;
        }

        m_pool->deallocate(p);
    }

    template<typename U>
    bool operator==(const PoolAllocator<U, Pool>& other) const noexcept
    {
        return m_pool == other.m_pool;
    }

    template<typename U>
    bool operator!=(const PoolAllocator<U, Pool>& other) const noexcept
    {
        return !(*this == other);
    }

private:
    template<typename U, typename OtherPool>
    friend class PoolAllocator;

    Pool* m_pool;
};

}
#endif // TOOLS_MEMORY_POOL_H
